import sys
import math

import gym
import d4rl
import torch

import numpy as np

from model import *
from utils import *

from train_fun import *

# 获取设备
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 生成一个环境对象
env=gym.make('hammer-human-v0')
# 获取dataset对象
dataset=env.get_dataset()

# 分割d4rl汇总的示范数据
obs_segs,act_segs=traj_breaks(dataset)
# list转换为np.array类型
obs_arr=get_array(obs_segs)
act_arr=get_array(act_segs)

# 生成torch.Dataset对象
torch_dataset=get_dataSet(obs_arr,act_arr,device=device)
# 生成torch.DataLoader对象
torch_loader=get_dataLoader(torch_dataset)

# 初始化模型对象的参数设定
input_dim=46
hidden_dim=64
latent_dim=32
act_dim=26
# 生成SOAS类型对象
soas=SOAS(input_dim,hidden_dim,latent_dim,act_dim,act_layer_num=5).to(device)
# 生成优化器对象
optimizer=torch.optim.Adam(soas.parameters(),lr=0.0001)

# 迭代的轮数
iteration_num=3000
# 关于mask01的标准差的设定
sigma_upper=0.51
sigma_lower=0.38
sigma_decay=(sigma_upper-sigma_lower)/(iteration_num)
sigma_epoch=sigma_upper
# 真正执行训练的循环
for epoch in range(1,iteration_num+1):
    train_with_latentKL_mask01(epoch,soas,torch_loader,optimizer,sigma_epoch,mse_loss_coef=0.0,latent_loss_coef=0.0,mask01_loss_coef=1.0)
